// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>

#include <nebula/bits/bit_util.h>
#include <nebula/bits/bitmap_reader.h>

namespace nebula::internal {

    // A function that visits each bit in a bitmap and calls a visitor function with a
    // boolean representation of that bit. This is intended to be analogous to
    // GenerateBits.
    template<class Visitor>
    void VisitBits(const uint8_t *bitmap, int64_t start_offset, int64_t length,
                   Visitor &&visit) {
        BitmapReader reader(bitmap, start_offset, length);
        for (int64_t index = 0; index < length; ++index) {
            visit(reader.IsSet());
            reader.next();
        }
    }

    // Like VisitBits(), but unrolls its main loop for better performance.
    template<class Visitor>
    void VisitBitsUnrolled(const uint8_t *bitmap, int64_t start_offset, int64_t length,
                           Visitor &&visit) {
        if (length == 0) {
            return;
        }

        // Start by visiting any bits preceding the first full byte.
        int64_t num_bits_before_full_bytes =
                bit_util::RoundUpToMultipleOf8(start_offset) - start_offset;
        // Truncate num_bits_before_full_bytes if it is greater than length.
        if (num_bits_before_full_bytes > length) {
            num_bits_before_full_bytes = length;
        }
        // Use the non loop-unrolled VisitBits since we don't want to add branches
        VisitBits<Visitor>(bitmap, start_offset, num_bits_before_full_bytes, visit);

        // Shift the start pointer to the first full byte and compute the
        // number of full bytes to be read.
        const uint8_t *first_full_byte = bitmap + bit_util::CeilDiv(start_offset, 8);
        const int64_t num_full_bytes = (length - num_bits_before_full_bytes) / 8;

        // Iterate over each full byte of the input bitmap and call the visitor in
        // a loop-unrolled manner.
        for (int64_t byte_index = 0; byte_index < num_full_bytes; ++byte_index) {
            // Get the current bit-packed byte value from the bitmap.
            const uint8_t byte = *(first_full_byte + byte_index);

            // execute the visitor function on each bit of the current byte.
            visit(bit_util::GetBitFromByte(byte, 0));
            visit(bit_util::GetBitFromByte(byte, 1));
            visit(bit_util::GetBitFromByte(byte, 2));
            visit(bit_util::GetBitFromByte(byte, 3));
            visit(bit_util::GetBitFromByte(byte, 4));
            visit(bit_util::GetBitFromByte(byte, 5));
            visit(bit_util::GetBitFromByte(byte, 6));
            visit(bit_util::GetBitFromByte(byte, 7));
        }

        // write any leftover bits in the last byte.
        const int64_t num_bits_after_full_bytes = (length - num_bits_before_full_bytes) % 8;
        VisitBits<Visitor>(first_full_byte + num_full_bytes, 0, num_bits_after_full_bytes,
                           visit);
    }

}  // namespace nebula::internal
